import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  
import os
import output_function as of

def laplace_5(myfun,dx,dy):
    """
    计算5点差分的拉普拉斯算子。
    :param myfun: 二维数组，要计算拉普拉斯算子的函数。
    :param dx: x方向的网格大小。
    :param dy: y方向的网格大小。
    :return: 拉普拉斯算子的结果。
    """
    diff_y = (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x+diff_y

def convection_operator(myfun,u,v,dx,dy):
    """
    计算对流算子。
    :param myfun: 二维数组，要计算对流算子的函数。
    :param u: 二维数组，x方向的速度。
    :param v: 二维数组，y方向的速度。
    :param dx: x方向的网格大小。
    :param dy: y方向的网格大小。
    :return: 对流算子的结果。
    """
    """x向对流速度u造成的输运,u \partial_x f"""
    uw_x = 0.5*u[1:-1,1:-1]  * (myfun[1:-1, 2:] - myfun[1:-1, :-2])/ dx
    """v \partial_y f"""
    vw_y = 0.5*v[1:-1,1:-1]  * (myfun[2:,1:-1] - myfun[:-2,1:-1])/ dy

    return uw_x,vw_y

class Simulation:
    def __init__(self, config):
        self.dx = config['dx']
        self.dy = config['dy']
        self.Lx = config['Lx']
        self.Ly = config['Ly']
        self.Nx = config['Nx']
        self.Ny = config['Ny']
        self.prandtl = config['prandtl']
        self.Ra = config['Ra']
        self.dt = config['dt']
        self.ntime = config['ntime']
        self.ndiag = config['ndiag']
        self.n_down=2   #底边密度
        self.n_up=1

    def initial_conditions_0(self):
        """初始化，从零开始"""
        self.x = np.linspace(0, Lx, Nx)
        self.y = np.linspace(0, Ly, Ny)
        self.xx, self.yy = np.meshgrid(self.x, self.y)  
        """设置初始流函数场"""
        self.psi=np.zeros_like(self.yy)
        self.u = np.zeros_like(self.xx)
        self.v = np.zeros_like(self.yy)
        """设置均匀密度场"""
        n_init0_y=np.linspace(self.n_down,self.n_up,Ny) #线性，因为扩散
        """矩阵形式,一个个赋值.第一次用到指标，j在前，i在后"""
        n_init0_matrix=np.zeros_like(self.yy)   
        for i in range(0,Nx):
                for j in range(0,Ny):
                    n_init0_matrix[j][i]=n_init0_y[j]

        """施加密度随机扰动"""
        n_density=np.copy(n_init0_matrix)
        # 设定噪声的标准差（std_dev），根据需要调整
        std_dev = 0.001  # 这是一个较小的噪声幅度
        #noise = np.random.normal(scale=std_dev, size=n_density.shape)
        noise=np.random.uniform(-std_dev, std_dev, size=n_density.shape)
        for i in range(0,Nx):
            for j in range(0,Ny*2//3):
                n_density[j][i]=n_density[j][i]*abs(1+noise[j][i])
        self.n_density=np.copy(n_density)

        """涡度场初始化"""
        self.omega = np.zeros_like(self.xx)
    
    def run_0(self):
        self.initial_conditions_0()
        self.out_iter = 0  
        while self.out_iter < self.ntime:
            self.out_iter += 1



    

# 给定的网格步长
dx = 0.01  # 例如，我们希望每个网格点之间的距离是 0.01
dy = 0.01  # 同上

# 给定的模拟区域长度
Lx = 1.0  # 模拟区域在 x 方向的长度
Ly = 1.0  # 模拟区域在 y 方向的长度

# 计算网格点数量
Nx = int(Lx / dx) + 1  # 加 1 是为了确保网格点覆盖整个模拟区域，包括边界
Ny = int(Ly / dy) + 1  # 同上

# 打印结果
print(f"Number of grid points in x-direction (Nx): {Nx}")
print(f"Number of grid points in y-direction (Ny): {Ny}")

dx=Lx/(Nx-1)
dy=Ly/(Ny-1)

prandtl=0.1 #Pr,普朗特数
Ra=200000 #Ra*,瑞利数

dt = 0.00002  #时间步长
ntime=1000  #时间步数
ndiag=10 #每隔多少时间输出一次

myconfig = {
    'dx': dx,
    'dy': dy,
    'Lx': Lx,
    'Ly': Ly,
    'Nx': Nx,
    'Ny': Ny,
    'prandtl':prandtl,
    'Ra':Ra,
    'dt': dt,
    'ntime': ntime,
    'ndiag': ndiag
}

# 使用myconfig字典来实例化Simulation类
sim = Simulation(myconfig)
sim.run_0()